#!/usr/bin/env python3
"""
ZMQ测试发布者，用于向127.0.0.1:8000发送KV缓存事件测试数据
"""
import time
import logging
import zmq
from msgspec.msgpack import Encoder
from src.zmq_manager import KVEventBatch, BlockStored, BlockRemoved, AllBlocksCleared


# 配置测试用的logger
logging.basicConfig(
    level=logging.INFO,
    format='%(asctime)s [%(levelname)s] %(name)s: %(message)s',
    datefmt='%Y-%m-%d %H:%M:%S'
)
test_logger = logging.getLogger("zmq_test")


def create_test_events():
    """创建测试事件数据"""
    events = []
    
    # 测试事件1: 存储一个新的块序列
    block_stored_1 = BlockStored(
        block_hashes=[1001, 1002, 1003],
        parent_block_hash=None,  # 第一个块没有父块
        token_ids=[1, 2, 3, 4, 5, 6, 7, 8, 9],  # 9个token，block_size=3时会分成3块
        block_size=3,
        lora_id=None
    )
    events.append(block_stored_1)
    
    # 测试事件2: 在现有块基础上继续存储
    block_stored_2 = BlockStored(
        block_hashes=[1004, 1005],
        parent_block_hash=1003,  # 基于前面的块继续
        token_ids=[10, 11, 12, 13, 14, 15],  # 6个token，block_size=3时会分成2块
        block_size=3,
        lora_id=None
    )
    events.append(block_stored_2)
    
    # 测试事件3: 另一个独立的块序列
    block_stored_3 = BlockStored(
        block_hashes=[2001],
        parent_block_hash=None,  # 新的独立序列
        token_ids=[100, 101, 102],
        block_size=3,
        lora_id=1  # 带lora_id
    )
    events.append(block_stored_3)
    
    return events


def create_removal_events():
    """创建删除事件数据"""
    return [
        BlockRemoved(block_hashes=[1004, 1005]),  # 删除一些块
    ]


def create_clear_event():
    """创建清空事件"""
    return [AllBlocksCleared()]


def main():
    test_logger.info("启动ZMQ测试发布者...")
    
    # 创建ZMQ上下文和发布者socket
    context = zmq.Context()
    publisher = context.socket(zmq.PUB)
    
    try:
        # 绑定到127.0.0.1:8000
        publisher.bind("tcp://127.0.0.1:8000")
        test_logger.info("ZMQ发布者已绑定到 tcp://127.0.0.1:8000")
        
        # 等待一下让连接建立
        time.sleep(1)
        
        encoder = Encoder()
        
        # 发送测试数据
        test_logger.info("=== 发送存储事件测试数据 ===")
        test_events = create_test_events()
        for i, event in enumerate(test_events, 1):
            event_batch = KVEventBatch(ts=time.time(), events=[event])
            payload = encoder.encode(event_batch)
            
            # 发送多部分消息: [topic, subtopic, payload]
            publisher.send_multipart([
                b"kv-events",  # topic
                b"",           # subtopic (空)
                payload        # 实际数据
            ])
            
            test_logger.info("已发送存储事件 %d: %s", i, type(event).__name__)
            if isinstance(event, BlockStored):
                test_logger.info("  - 块哈希: %s", event.block_hashes)
                test_logger.info("  - 父块哈希: %s", event.parent_block_hash)
                test_logger.info("  - Token数量: %d", len(event.token_ids))
                test_logger.info("  - 块大小: %d", event.block_size)
            
            time.sleep(0.5)  # 间隔发送
        
        test_logger.info("=== 等待5秒后发送删除事件 ===")
        time.sleep(5)
        
        removal_events = create_removal_events()
        for i, event in enumerate(removal_events, 1):
            event_batch = KVEventBatch(ts=time.time(), events=[event])
            payload = encoder.encode(event_batch)
            
            publisher.send_multipart([
                b"kv-events",
                b"",
                payload
            ])
            
            test_logger.info("已发送删除事件 %d: %s", i, type(event).__name__)
            test_logger.info("  - 删除块哈希: %s", event.block_hashes)
            time.sleep(0.5)
        
        test_logger.info("=== 等待5秒后发送清空事件 ===")
        time.sleep(5)
        
        clear_events = create_clear_event()
        for i, event in enumerate(clear_events, 1):
            event_batch = KVEventBatch(ts=time.time(), events=[event])
            payload = encoder.encode(event_batch)
            
            publisher.send_multipart([
                b"kv-events",
                b"",
                payload
            ])
            
            test_logger.info("已发送清空事件 %d: %s", i, type(event).__name__)
            time.sleep(0.5)
        
        test_logger.info("=== 继续发送更多测试数据 ===")
        # 发送一些用于匹配测试的数据
        match_test_event = BlockStored(
            block_hashes=[3001, 3002],
            parent_block_hash=None,
            token_ids=[1, 2, 3, 4, 5, 6],  # 与第一个事件有重叠，用于测试匹配
            block_size=3,
            lora_id=None
        )
        
        event_batch = KVEventBatch(ts=time.time(), events=[match_test_event])
        payload = encoder.encode(event_batch)
        
        publisher.send_multipart([
            b"kv-events",
            b"",
            payload
        ])
        
        test_logger.info("已发送匹配测试事件:")
        test_logger.info("  - 块哈希: %s", match_test_event.block_hashes)
        test_logger.info("  - Token: %s", match_test_event.token_ids)
        
        test_logger.info("=== 测试数据发送完成，保持发布者运行... ===")
        test_logger.info("按 Ctrl+C 停止发布者")
        
        # 保持运行
        while True:
            time.sleep(1)
            
    except KeyboardInterrupt:
        test_logger.warning("收到中断信号，停止发布者...")
    except Exception as e:
        test_logger.error("发布者错误: %s", e)
    finally:
        publisher.close()
        context.term()
        test_logger.info("ZMQ发布者已停止")


if __name__ == "__main__":
    main()